"""
基于LangGraph的调度推荐工作流
"""
from langgraph.graph import StateGraph, END
from langgraph.prebuilt import ToolNode
from typing import TypedDict, Annotated, List, Dict, Any
import json
import os
import sys

# 添加项目路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(__file__))))

from ai.agents.tools.data_collection_tool import data_collection_tool
from ai.agents.tools.similarity_calculation_tool import similarity_calculation_tool
from ai.agents.tools.area_recommendation_tool import area_recommendation_tool

# 状态定义
class DispatchState(TypedDict):
    driver_id: int
    time_category: str
    data_file: str
    similar_drivers: List[Dict]
    recommendations: List[Dict]
    error: str
    success: bool
    actual_time_category: str

# 数据收集节点
def collect_data(state: DispatchState) -> DispatchState:
    """数据收集节点"""
    try:
        driver_id = state["driver_id"]
        time_category = state["time_category"]
        
        # 构建查询参数
        data_query = json.dumps({
            "days": 7,
            "output_file": f"driver_data_{driver_id}_{time_category}.csv",
            "time_category": time_category
        })
        
        # 调用数据收集工具
        data_result = data_collection_tool.invoke(data_query)
        
        # 解析结果
        if "数据收集完成：" in data_result:
            data_json_str = data_result.split("数据收集完成：")[1]
            data_stats = json.loads(data_json_str)
            actual_time_category = data_stats.get('query_params', {}).get('actual_time_category', time_category)
            
            # 使用实际生成的文件名
            actual_file_path = data_stats.get('file_path', '')
            if actual_file_path:
                # 从文件路径中提取文件名
                import os
                data_file = os.path.basename(actual_file_path)
            else:
                # 回退到默认文件名
                data_file = f"driver_data_{driver_id}_{actual_time_category}.csv"
            
            return {
                **state,
                "data_file": data_file,
                "actual_time_category": actual_time_category,
                "success": True,
                "error": ""
            }
        else:
            # 提取具体的错误信息
            error_message = data_result
            if "数据收集失败：" in data_result:
                error_message = data_result.split("数据收集失败：")[1]
            return {
                **state,
                "success": False,
                "error": f"数据收集失败: {error_message}"
            }
            
    except Exception as e:
        return {
            **state,
            "success": False,
            "error": f"数据收集异常: {str(e)}"
        }

# 相似度计算节点
def calculate_similarity(state: DispatchState) -> DispatchState:
    """相似度计算节点"""
    try:
        if not state["success"]:
            return state
            
        driver_id = state["driver_id"]
        data_file = state["data_file"]
        
        # 构建查询参数
        similarity_query = json.dumps({
            "input_file": data_file,
            "target_driver_id": driver_id,
            "top_k": 10
        })
        
        # 调用相似度计算工具
        similarity_result = similarity_calculation_tool.invoke(similarity_query)
        
        # 解析结果
        if "相似度计算完成：" in similarity_result:
            similarity_json_str = similarity_result.split("相似度计算完成：")[1]
            similarity_data = json.loads(similarity_json_str)
            similar_drivers = similarity_data.get("similar_drivers", [])
            
            return {
                **state,
                "similar_drivers": similar_drivers,
                "success": True,
                "error": ""
            }
        else:
            # 提取具体的错误信息
            error_message = similarity_result
            if "相似度计算失败：" in similarity_result:
                error_message = similarity_result.split("相似度计算失败：")[1]
            return {
                **state,
                "success": False,
                "error": f"相似度计算失败: {error_message}"
            }
            
    except Exception as e:
        return {
            **state,
            "success": False,
            "error": f"相似度计算异常: {str(e)}"
        }

# 区域推荐节点
def recommend_areas(state: DispatchState) -> DispatchState:
    """区域推荐节点"""
    try:
        if not state["success"]:
            return state
            
        data_file = state["data_file"]
        similar_drivers = state["similar_drivers"]
        actual_time_category = state["actual_time_category"]
        
        # 构建查询参数
        recommendation_query = json.dumps({
            "input_file": data_file,
            "similar_drivers": similar_drivers,
            "time_category": actual_time_category,
            "top_areas": 5
        })
        
        # 调用区域推荐工具
        recommendation_result = area_recommendation_tool.invoke(recommendation_query)
        
        # 解析结果
        if "区域推荐完成：" in recommendation_result:
            recommendation_json_str = recommendation_result.split("区域推荐完成：")[1]
            recommendation_data = json.loads(recommendation_json_str)
            recommendations = recommendation_data.get("top_recommendations", [])
            
            return {
                **state,
                "recommendations": recommendations,
                "success": True,
                "error": ""
            }
        else:
            return {
                **state,
                "success": False,
                "error": "区域推荐失败"
            }
            
    except Exception as e:
        return {
            **state,
            "success": False,
            "error": f"区域推荐异常: {str(e)}"
        }

# 构建调度图
def create_dispatch_graph():
    """创建调度推荐工作流图"""
    workflow = StateGraph(DispatchState)
    
    # 添加节点
    workflow.add_node("collect_data", collect_data)
    workflow.add_node("calculate_similarity", calculate_similarity)
    workflow.add_node("recommend_areas", recommend_areas)
    
    # 设置边
    workflow.set_entry_point("collect_data")
    workflow.add_edge("collect_data", "calculate_similarity")
    workflow.add_edge("calculate_similarity", "recommend_areas")
    workflow.add_edge("recommend_areas", END)
    
    return workflow.compile()

# 热点分析图
def create_hotspot_graph():
    """创建热点分析工作流图"""
    workflow = StateGraph(DispatchState)
    
    # 热点分析只需要数据收集
    workflow.add_node("collect_data", collect_data)
    workflow.set_entry_point("collect_data")
    workflow.add_edge("collect_data", END)
    
    return workflow.compile() 